import torchvision
from torch.utils.data import DataLoader

from datasets.dataset import  UcfFramesDataSet
from tools.transforms import GroupMultiScaleCrop, GroupRandomHorizontalFlip, Stack, ToTorchFormatTensor, GroupNormalize

if __name__ == '__main__':

    train_dataset = UcfFramesDataSet(
        ann_file=r"data\ucf101\ucf101_train_split_1_rawframes.txt",
        t_num=3,
        test_mode=True,
        transform=torchvision.transforms.Compose([
            torchvision.transforms.Compose([GroupMultiScaleCrop(224, [1, .875, .75, .66]),
                                            GroupRandomHorizontalFlip(is_flow=False)]),
            Stack(roll=False),
            ToTorchFormatTensor(div=True),
            GroupNormalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    )
    #

    train_loader = DataLoader(train_dataset, shuffle=False,
                              num_workers=2, pin_memory=True,
                              batch_size=2,
                              # collate_fn=yolo_dataset_collate
                              )

    for i, (input, target) in enumerate(train_loader):
        print(input.shape)

    pass
